'''
-*- coding: utf-8 -*-
@File  : 量化.py
@Author: Shanmh
@Time  : 2024/03/20 上午10:15
@Function：
'''

import os
import openvino as ov



def export_int8(pt_path):
    path_,name_=os.path.split(pt_path)
    save_path=os.path.join(path_,name_[:-3]+"_openvino_model")
    # 定义模型路径和文件名
    model_xml = f'{save_path}/{name_[:-3]}.xml'
    model_bin = f'{save_path}/{name_[:-3]}.bin'

    # 定义校准数据集路径
    calibration_dataset = '/path/to/calibration/dataset'

    # 定义输出目录
    output_dir = 'quantized_model'

    # 创建 OpenVINO 的 Inference Engine 核心对象
    ie = ov.Core()

    # 加载浮点模型
    net = ie.read_network(model=model_xml, weights=model_bin)

    # 指定量化的数据类型为 INT8
    net.precision = 'INT8'

    # 创建 INT8 量化参数
    int8_config = {
        'calibration_dataset': calibration_dataset,
        'output_dir': output_dir,
        'input_shape': [1, 3, height, width],
        'mean_values': [mean_values],
        'scale_values': [scale_values]
    }

    # 进行模型量化
    quantized_net = ie.quantize(net, int8_config)

    # 定义量化后的模型文件路径和文件名
    quantized_model_xml = os.path.join(output_dir, 'model_int8.xml')
    quantized_model_bin = os.path.join(output_dir, 'model_int8.bin')

    # 保存量化后的模型
    quantized_net.save_model(quantized_model_xml, quantized_model_bin)